import akshare as ak
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime, timedelta


# 1. 数据获取模块
def fetch_industry_data():
    """获取行业板块实时数据"""
    # 获取所有行业板块列表
    industry_list = ak.stock_board_industry_name_em()["板块名称"].tolist()

    # 获取板块历史数据（近20日）
    end_date = datetime.now().strftime("%Y%m%d")
    start_date = (datetime.now() - timedelta(days=30)).strftime("%Y%m%d")
    industry_data = {}

    for industry in industry_list[:10]:  # 示例仅取前10个板块加速计算
        try:
            df = ak.stock_board_industry_hist_em(
                symbol=industry, start_date=start_date, end_date=end_date
            )
            df["板块名称"] = industry
            industry_data[industry] = df
        except:
            continue

    return pd.concat(industry_data.values())


# 2. 板块强度计算模块
def calculate_strength(df):
    """计算板块强度指标（基于成交量均线轮动+市值加权）[2,5](@ref)"""
    # 计算成交量均线
    df["MA5_VOL"] = df["成交量"].rolling(5).mean()
    df["MA20_VOL"] = df["成交量"].rolling(20).mean()
    print(df.head())
    # 计算市值加权涨跌幅（假设用流通市值代替）
    df["加权涨跌幅"] = df["涨跌幅"] * df["成交额"]
    df["板块强度"] = df["加权涨跌幅"].rolling(5).mean() + (
        df["MA5_VOL"] / df["MA20_VOL"]
    )

    return df.groupby("板块名称", group_keys=False).apply(lambda x: x.iloc[-1])


# 3. 轮动信号生成模块
def generate_rotation_signal(strength_df):
    """生成轮动信号（TOP3强势板块）"""
    strength_df = strength_df.sort_values("板块强度", ascending=False)
    top_industries = strength_df.head(3)["板块名称"].tolist()

    # 信号逻辑：若板块强度连续3日上升则买入
    signals = {}
    for industry in top_industries:
        industry_data = strength_df[strength_df["板块名称"] == industry]
        if len(industry_data) >= 3:
            trend_up = all(
                industry_data["板块强度"].iloc[-3:]
                > industry_data["板块强度"].iloc[-4:-1]
            )
            signals[industry] = "买入" if trend_up else "观望"
    return signals


# 4. 可视化模块
def plot_rotation(strength_df):
    """绘制板块强度热力图[7](@ref)"""
    pivot_df = strength_df.pivot(
        index="日期", columns="板块名称", values="板块强度"
    ).fillna(0)
    plt.figure(figsize=(15, 8))
    # sns.heatmap(pivot_df.T, cmap="RdYlGn", annot=False)
    plt.title("板块强度轮动热力图（红色=强势）")
    plt.savefig("板块轮动热力图.png")


# 5. 主流程执行
if __name__ == "__main__":
    # 获取数据
    industry_df = fetch_industry_data()
    print("【行业板块数据】=", industry_df.head())

    # 计算板块强度
    strength_df = calculate_strength(industry_df)
    print("【板块强度数据】=", strength_df.head())

    # 生成信号
    signals = generate_rotation_signal(strength_df)
    print("【轮动信号】", signals)

    # 可视化
    # plot_rotation(strength_df[-30:])  # 仅展示最近30天数据

    # 输出最强板块龙头股（需补充个股数据接口）
    # 示例：获取光伏板块龙头（实际需结合个股强度筛选）
    if "光伏设备" in signals:
        dragon_stock = ak.stock_board_industry_cons_em(symbol="光伏设备").loc[0]["代码"]
        print(f"光伏设备板块龙头股：{dragon_stock}")
